import tensorflow as tf
from tensorflow.keras.layers import Embedding, Input, add, BatchNormalization, dot
from tensorflow.keras.models import Model

from GCN_layers import GraphConvolution, AsymmetricGraphConvolution, DynamicGraphConvolution
from bpr_loss import BPRLossWithoutSample
from cross_gate import CrossGate
from tensor_utils import to_dense_tensor

custom_objects = {
    'GraphConvolution': GraphConvolution,
    'AsymmetricGraphConvolution': AsymmetricGraphConvolution,
    'DynamicGraphConvolution': DynamicGraphConvolution,
    'CrossGate': CrossGate,
    'BPRLossWithoutSample': BPRLossWithoutSample
}


def build_gcnrec_model(num_users, num_items,
                       graph_uu, graph_ii, graph_ui, item_features,
                       emb_dim, graph_dim, bpr_loss_lambda):
    graph_ui = to_dense_tensor(graph_ui, dtype=tf.float32)

    user_input = Input(shape=(num_users,), dtype='int32', name='user_input')
    item_input = Input(shape=(num_items,), dtype='int32', name='item_input')

    user_emb = Embedding(input_dim=num_users, output_dim=emb_dim, name='user_embedding',
                         embeddings_initializer='random_normal', input_length=1)(user_input)
    item_emb = Embedding(input_dim=num_items, output_dim=emb_dim, name='item_embedding',
                         embeddings_initializer='random_normal', input_length=1)(item_input)

    user_static = GraphConvolution(units=graph_dim, graph=graph_uu,
                                   activation='relu', name='user_static_graph')(user_emb)
    user_dynamic = DynamicGraphConvolution(units=graph_dim, graph=graph_uu,
                                           activation='relu', name='user_dynamic_graph')(user_emb)
    user_inter = AsymmetricGraphConvolution(units=graph_dim, graph=graph_ui,
                                            activation='relu', name='user_inter_graph')(item_emb)

    user_agg = add([user_static, user_dynamic, user_inter])
    user_agg = BatchNormalization(axis=2, name='user_out')(user_agg)

    item_static = GraphConvolution(units=graph_dim, graph=graph_ii,
                                   activation='relu', name='item_static_graph')(item_emb)
    item_dynamic = DynamicGraphConvolution(units=graph_dim, graph=graph_ii,
                                           activation='relu', name='item_dynamic_graph')(item_emb)
    item_inter = AsymmetricGraphConvolution(units=graph_dim, graph=tf.transpose(graph_ui),
                                            activation='relu', name='item_inter_graph')(user_emb)

    item_agg = add([item_static, item_dynamic, item_inter])
    item_agg = BatchNormalization(axis=2)(item_agg)
    item_agg = CrossGate(units=graph_dim, features=item_features, activation='sigmoid',
                         out_activation='sigmoid', name='item_cross_gate')(item_agg)
    item_agg = BatchNormalization(axis=2, name='item_out')(item_agg)

    rating = dot([user_agg, item_agg], axes=(2, 2))

    pos_item_id_input = Input(shape=(num_users,), batch_size=1, dtype='int32', name='pos_item_id_input')
    neg_item_id_input = Input(shape=(num_users,), batch_size=1, dtype='int32', name='neg_item_id_input')

    # loss = bpr_loss(lamb=bpr_loss_lambda)(user_agg, item_agg, pos_item_id_input, neg_item_id_input)

    loss = BPRLossWithoutSample(lamb=bpr_loss_lambda, name='loss')(
        [user_agg, item_agg, pos_item_id_input, neg_item_id_input])

    gcnrec_model = Model(inputs=[user_input, item_input, pos_item_id_input, neg_item_id_input],
                         outputs=[rating])

    gcnrec_model.add_loss(loss)

    return gcnrec_model
